package cn.seecoder.ai.algorithm.ml.classification
import cn.seecoder.ai.algorithm.ml.BaseClassification
import cn.seecoder.ai.model.bo.TrainParamsBO
import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.classification.MultilayerPerceptronClassifier


class MultilayerPerceptronClassification extends BaseClassification {


  override def buildMachineLearningStage(trainParams: TrainParamsBO): PipelineStage = {

    // 三层感知机：全感知机
    val layers = Array[Int](
      trainParams.featureDimensionNum,
      trainParams.featureDimensionNum,
      trainParams.labelUniqueNum)


    val mp = new MultilayerPerceptronClassifier()
      .setFeaturesCol("indexedFeatures")
      .setLabelCol("label")
      .setSeed(trainParams.seed)
      .setLayers(layers)
      .setBlockSize(trainParams.blockSize)
      .setMaxIter(trainParams.maxIter)
    mp

  }

}
